package logic;

import java.util.ArrayList;
import java.util.Enumeration;
import java.util.HashMap;
import weka.classifiers.trees.Id3;
import weka.core.Attribute;
import weka.core.Instance;
import weka.core.Instances;

public class MushroomNB extends MushroomClassifier {
    public ArrayList<HashMap<Character, HashMap<Character, Integer>>> counts;
    
    public MushroomNB() {
        algo = "NB";
    }
    
    public final void initCounts() {
        counts = new ArrayList<HashMap<Character, HashMap<Character, Integer>>>();
        
        Attribute clsAtr = dataSet.classAttribute();
        for (int i = 0; i < dataSet.numAttributes(); i++) {
            HashMap<Character, HashMap<Character, Integer>> attrCounts = new HashMap<Character, HashMap<Character, Integer>>();
            Attribute atr = dataSet.attribute(i);
            for (int j = 0; j < atr.numValues(); j++) {
                HashMap<Character, Integer> valCounts = new HashMap<Character, Integer>();
                for (int k = 0; k < clsAtr.numValues(); k++) {
                    valCounts.put(clsAtr.value(k).charAt(0), 0);
                }
                attrCounts.put(atr.value(j).charAt(0), valCounts);
            }
            counts.add(attrCounts);
        }
    }
    
    public void addCount(int attrIdx, char attrValue, char classValue) {
        HashMap<Character, Integer> count = counts.get(attrIdx).get(attrValue);
        int inc = count.get(classValue) + 1;
        count.put(classValue, inc);
    }
    
    @Override
    public void buildClassifier(Instances is) throws Exception {
        dataSet = new Instances(is);
        initCounts();
        
        for (int i = 0; i < is.numInstances(); i++) {
            Instance in = is.instance(i);
            String[] strAttr = in.toString().split(",");
            char cls = strAttr[strAttr.length-1].charAt(0);
            for (int j = 0; j < is.numAttributes(); j++) {
                addCount(j, strAttr[j].charAt(0), cls);
            }
        }
    }
    
    @Override
    public char classify(Instance i) {
        HashMap<Character, Double> clsProb = new HashMap<Character, Double>();
        int numInst = dataSet.numInstances();
        Attribute clsAttr = dataSet.classAttribute();
        
        Instance in = initTestInstance(i);
        
        for (int j = 0; j < clsAttr.numValues(); j++) {
            char cls = clsAttr.value(j).charAt(0);
            double num = counts.get(dataSet.classIndex()).get(cls).get(cls);
            double d = num/numInst;
            
            String[] spl = in.toString().split(",");
            for (int k = 0; k < spl.length - 1; k++) {
                char value = spl[k].charAt(0);
                int numAttr = counts.get(k).get(value).get(cls);
                
                d = (d*numAttr)/num;
            }
            clsProb.put(cls, d);
        }
        
        double max = 0;
        char classifiedAs = '\0';
        for (char cls : clsProb.keySet()) {
            if (clsProb.get(cls) > max) {
                max = clsProb.get(cls);
                classifiedAs = cls;
            }
        }
        
        return classifiedAs;
    }
    
    @Override
    public void learn(Instance i) {
        dataSet.add(new Instance(i));
        
        char label = classify(i);
        i.setClassValue(label+"");
        String[] spl = i.toString().split(",");
        for (int j = 0; j < spl.length; j++) {
            addCount(j, spl[j].charAt(0), label);
        }
    }
    
    public int getNumInstance() {
        int numInst = 0;
        
        for (char cls : MushroomInfo.CLASS_VALUES_SHORT) {
            numInst += counts.get(MushroomInfo.CLASS_IDX).get(cls).get(cls);
        }
        
        return numInst;
    }
    
    public String globalInfo() {
        return "Naive Bayes untuk Mushroom Dataset";
    }
    
    public String stringModel() {
        ArrayList<HashMap<Character, String>> longs = MushroomInfo.singleton.ATTRIBUTE_VALUES_LONG;
        StringBuilder s = new StringBuilder();
        
        int i = 0;
        for (HashMap<Character, HashMap<Character, Integer>> attrCounts : counts) {
            if (attrCounts == null)
                continue;
            
            s.append(MushroomInfo.ATTRIBUTE_NAMES[i]).append("\n");
            for (char attr : attrCounts.keySet()) {
                s.append(longs.get(i).get(attr)).append("    ").append(attrCounts.get(attr).toString()).append("\n");
            }
            s.append("\n");
            i++;
        }
        
        return s.toString();
    }
    
    public static void main(String[] args) {
        Instances inst = Loader.load("D:\\mushroom.arff");
        inst = Loader.preProcess(inst, Loader.UNSUPERVISED_REPLACE_MISSING_VALUES);
//        MushroomNB nb = (MushroomNB)Loader.loadModel("D:\\mushroomNB.model");
        MushroomNB nb = new MushroomNB();
        try {
            nb.buildClassifier(inst);
        } catch (Exception e) {
            e.printStackTrace();
        }
        
        System.out.println(nb.stringModel());
        Instance f = inst.firstInstance();
        System.out.println(nb.classifyInstance(f));
        System.out.println(f.classValue());
        System.out.println(nb.accuracy(inst));
        
        Loader.saveModel(nb, Loader.MODEL_MUSHROOM_NB, "D:\\mushroomNB.model");
    }
}
